import json
import os
import time

#定义要修改的字段名称（目前仅支持指定2项参数的情况）
field_names = ['thread_count', 'last_layer_nprobe']
field_values = [
    # [1, list(range(40,41,20))],
    # [2, list(range(40,41,20))],
    # [4, list(range(40,41,20))],
    [8, list(range(40,41,20))],
    # [16, list(range(40,41,20))],
    # [32, list(range(40,41,20))],
    # [64, list(range(40,41,20))],
]
exec_file = './annlite_test_uint8'
# exec_file = './annlite_test_float32'

#定义其他字段的值
values = {
    "index_root_path": "../annlite_builder/bigann_1M_annlite_c32000_re2_pq32_R30_l1s0.1/",
    "max_query": 10000,
    "thread_count": 8,
    "recall_k": 10,
    "real_k": 10,
    "calculate_recall_of_layers": 0,
    "l2_io_limit": 10,
    "l2_update_ratio": 0.6,
    "l2_beamwidth_start": 8,
    "beamwidth_increase_factor": 1,
    "pipesearch_on": 1,
    "last_layer_nprobe": 20,
    "l3_graph_load_in_batch": 1,
    "bloomfilter_size": 16384,
    "bloomfilter_hash_count": 10,
}

#获取所有参数组合
def get_all_combinations(field_values):
    combinations = []
    for line in field_values:
        for value in line[1]:
            combinations.append((line[0], value))
    return combinations

value_sets = get_all_combinations(field_values)
#把所有参数组合以及参数的名字(field_names)打印成表格
print("All Test Values:")
print(field_names)
for value_set in value_sets:
    print(value_set)

#定义要修改的json文件路径
json_file_path = 'user_config.json'
#定义运行完成后的检索结果输出路径
result_file_path = 'result_detail_sum.csv'
#定义本自动测试脚本输出的整合结果文件路径
auto_tester_result_file_path = 'auto_tester_result.csv'
sleep_time = 0

if __name__ == "__main__":
    #首先加载user_config.json文件
    user_config = None
    with open(json_file_path, 'r') as f:
        user_config = json.load(f)

    #先将其他字段的值复制到user_config
    for key, value in values.items():
        user_config[key] = value

    start_time = time.time()

    #依次将字段值修改为列表中的值并运行
    first_loop = True
    for value_set in value_sets:
        for field_num in range(len(field_names)):
            user_config[field_names[field_num]] = value_set[field_num]
            print(f"set {field_names[field_num]} to {value_set[field_num]}")

        with open(json_file_path, 'w') as f:
            json.dump(user_config, f, indent=4)

        #运行make run
        print("start running")
        os.system(exec_file)
        time.sleep(sleep_time/2) #等待sleep_time秒等待结果文件写入

        #运行完成后打开result_file_path。
        with open(result_file_path, 'r') as f:
            results = f.readlines()
            #如果是第一次循环，直接复制到auto_tester_result_file_path；如果不是，则只复制第二行
            if first_loop:
                first_loop = False
                with open(auto_tester_result_file_path, 'w') as f:
                    #先写入两项可变参数的名称
                    for field_name in field_names:
                        f.write(field_name)
                        f.write(",")
                    f.write(results[0])
                    #写入可变参数的值
                    for field_num in range(len(field_names)):
                        f.write(str(value_set[field_num]))
                        f.write(",")
                    f.write(results[1])
            else:
                with open(auto_tester_result_file_path, 'a') as f:
                    f.write("\n")
                    #先写入可变参数的值
                    for field_num in range(len(field_names)):
                        f.write(str(value_set[field_num]))
                        f.write(",")
                    f.write(results[1])

        time.sleep(sleep_time) #等待sleep_time秒避免上一轮运行的影响

    end_time = time.time()
    print("Auto tester done, result write to", auto_tester_result_file_path)
    print("All Test Values:")
    for value_set in value_sets:
        print(value_set)
    print(f"Total time: {round(end_time - start_time, 2)} seconds")
    
